from typing import List, Dict, Any
from langchain_core.retrievers import BaseRetriever
from langchain_core.documents import Document
from langchain.vectorstores import VectorStore

class HybridRetriever(BaseRetriever):
    """先阈值过滤，再MMR的混合检索器"""
    vector_store:VectorStore
    threshold:float
    mmr_k:int
    lambda_mult:float
    
    def _get_relevant_documents(self, query: str, **kwargs) -> List[Document]:
        retriever_score_threshold = self.vector_store.as_retriever(
            search_type="similarity_score_threshold",  # 关键参数
            search_kwargs={
                "k": self.mmr_k * 3,                      # 最大返回数量（实际可能少于k）
                "score_threshold": self.threshold        # 最小相似度阈值（根据分数类型调整）
            }
        )
        # 第一阶段：阈值过滤
        docs_with_scores = retriever_score_threshold.invoke(query)
        filtered_docs = [doc for doc in docs_with_scores]
        if not docs_with_scores:
            return []  # 无达标文档时返回空
        # 第二阶段：MMR重排序
        mmr_docs = self.vector_store.max_marginal_relevance_search(
            query, 
            k=self.mmr_k,
            lambda_mult=self.lambda_mult,
            filter_docs=filtered_docs  # 仅在过滤后文档中运行MMR
        )
        return mmr_docs